package com.csx.governance.traffic.core.utils;

import com.csx.governance.traffic.core.Weighted;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.Iterator;
import java.util.List;
import java.util.Random;
import java.util.concurrent.ThreadLocalRandom;
import java.util.function.ToIntFunction;

/**
 * @author cuisongxu
 * @data 2023/4/2 17:04
 */
public class Weights {

    private Weights() {
    }

    /**
     * 使用ThreadLocalRandom的方式生成随机数的方式，按照权重比例获取元素
     */
    public static <T extends Weighted> T getByRandomWeight(List<T> weights) {
        return getByRandomWeight(weights, ThreadLocalRandom.current());
    }

    public static <T extends Weighted> T getByRandomWeight(List<T> weights, Random random) {
        return getByRandomWeight(weights, random::nextInt);
    }

    /**
     * 权重随机的方式获取一个元素
     *
     * @param originWeights        带权重的元素
     * @param integerToIntFunction 一个入参为权重总和的回调函数
     * @return 目标元素
     */
    @SuppressWarnings("unchecked")
    public static <T extends Weighted> T getByRandomWeight(List<T> originWeights, ToIntFunction<Integer> integerToIntFunction) {
        if (originWeights == null || originWeights.isEmpty()) {
            return null;
        }
        if (originWeights.size() == 1) {
            return originWeights.get(0);
        }
        List<T> weights = new ArrayList<>(originWeights);
        int totalWeight = 0;
        Iterator<T> it = weights.iterator();
        while (it.hasNext()) {
            T weighted = it.next();
            int weight = weighted.getWeight();
            if (weight <= 0) {
                it.remove();
                continue;
            }
            totalWeight += weight;
        }
        //说明集合元素的权重全都为0，则随机返回一个
        if (weights.isEmpty()) {
            int index = ThreadLocalRandom.current().nextInt(originWeights.size());
            return originWeights.get(index);
        }
        weights.sort(Comparator.comparingInt(Weighted::getWeight));
        int randomInt = integerToIntFunction.applyAsInt(totalWeight);
        Weighted weighted = weights.get(weights.size() - 1);
        int current = 0;
        for (Weighted weightedItem : weights) {
            int weight = weightedItem.getWeight();
            current += weight;
            int i = current - randomInt;
            if (i >= 0) {
                weighted = weightedItem;
                break;
            }
        }
        return (T) weighted;
    }
}
